import sys 
import os
import random
import torch
import numpy as np

sys.path.append(os.path.join(os.path.dirname(__file__), "../src"))
from lightning_attn_triton import lightning_attention
from lightning_attn_torch import reference_lightning_attention
from gen_test_input import gen_input, TEST_CASES, TEST_CASE_HEADER

sys.path.append(os.path.join(os.path.dirname(__file__), "../../../precision"))
from compare import check_operator_accuracy

sys.path.append(os.path.join(os.path.dirname(__file__), "../../../profiler"))
from triton_profiler import triton_profiler_wrapper, times_list_header

sys.path.append(os.path.join(os.path.dirname(__file__), "../../../utils"))
# from utils import write_csv
from templates import prof_template


# 运行设备
torch.npu.set_device(6)
# 随机种子
seed = 42
torch.manual_seed(seed)
np.random.seed(seed) 
random.seed(seed)



def prof_lightning_attention(
    test_case
):

    q, k, v, ed, kv_history, block_size = gen_input(*test_case)
    kv_history_clone = kv_history.clone()
    # profiling
    def call(): 
        # ref_output, ref_kv_cache, ref_kv = reference_lightning_attention(
        #     q, k, v, ed, kv_history, block_size)

        actual_output, actual_kv_cache, actual_kv = lightning_attention(
            q, k, v, ed, kv_history_clone, block_size)
        
    # do_bench
    times_list = triton_profiler_wrapper(call)
    print(times_list)
    return times_list


if __name__ == "__main__": 

    csv_header_list = TEST_CASE_HEADER + times_list_header
    prof_template("LightningAttn", TEST_CASES, prof_lightning_attention, csv_header_list)
